/*
 * Copyright 2023-2025, Haiku, Inc. All rights reserved.
 * Distributed under the terms of the MIT License.
 */

#include <arch_cpu_defs.h>
#include <boot/kernel_args.h>
#include <KernelExport.h>
#include <kernel.h>
#include <vm/vm.h>
#include <vm/vm_priv.h>
#include <vm/VMAddressSpace.h>

#include "LoongArchVMTranslationMap.h"

// LoongArchVMTranslationMap类实现

// 构造函数
LoongArchVMTranslationMap::LoongArchVMTranslationMap(bool kernel, phys_addr_t pageTable)
    : fPageTable(pageTable), fASID(0), fIsKernel(kernel)
{
    // 如果是内核映射，使用全局ASID 0
    if (kernel)
    {
        fASID = 0;
    }
    else
    {
        // 为用户空间分配一个新的ASID
        // 在实际实现中，应该有一个ASID分配器
        static uint32 nextASID = 1;
        fASID = nextASID++;
        if (nextASID >= 256) // 假设最多支持256个ASID
            nextASID = 1;
    }

    TRACE("LoongArchVMTranslationMap::ctor: %s, pageTable: %#" B_PRIxPHYSADDR ", ASID: %u\n",
          kernel ? "kernel" : "user", pageTable, fASID);
}

// 析构函数
LoongArchVMTranslationMap::~LoongArchVMTranslationMap()
{
    TRACE("LoongArchVMTranslationMap::dtor\n");
}

// 锁定映射
bool LoongArchVMTranslationMap::Lock()
{
    TRACE("LoongArchVMTranslationMap::Lock\n");
    return true;
}

// 解锁映射
void LoongArchVMTranslationMap::Unlock()
{
    TRACE("LoongArchVMTranslationMap::Unlock\n");
}

// 获取已映射大小
addr_t
LoongArchVMTranslationMap::MappedSize() const
{
    // 这个函数应该返回已映射的内存大小
    // 在实际实现中，应该遍历页表计算
    return 0;
}

// 计算映射所需的最大页数
size_t
LoongArchVMTranslationMap::MaxPagesNeededToMap(addr_t start, addr_t end) const
{
    // 计算映射指定范围所需的最大页表页数
    // 每个页表页可以映射多少内存取决于页表结构
    // 对于4级页表，每个页表页包含512个条目
    // 每个L3条目映射4KB，每个L2条目映射2MB，每个L1条目映射1GB

    // 简单实现：假设最坏情况，每个页需要所有级别的页表
    return ((end - start) / B_PAGE_SIZE + 1) * 4;
}

// 映射内存
status_t
LoongArchVMTranslationMap::Map(addr_t virtualAddress, phys_addr_t physicalAddress,
                               uint32 attributes, uint32 memoryType, vm_page_reservation *reservation)
{
    TRACE("LoongArchVMTranslationMap::Map: va %#" B_PRIxADDR ", pa %#" B_PRIxPHYSADDR ", attributes %#" B_PRIx32 "\n",
          virtualAddress, physicalAddress, attributes);

    // 检查地址对齐
    if ((virtualAddress & (B_PAGE_SIZE - 1)) != 0)
        return B_BAD_VALUE;
    if ((physicalAddress & (B_PAGE_SIZE - 1)) != 0)
        return B_BAD_VALUE;

    // 构建页表项标志
    uint64 flags = 0;

    // 设置访问权限
    if ((attributes & B_READ_AREA) != 0)
        flags |= kPteRead;
    if ((attributes & B_WRITE_AREA) != 0 || (attributes & B_KERNEL_WRITE_AREA) != 0)
        flags |= kPteWrite;
    if ((attributes & B_EXECUTE_AREA) != 0)
        flags |= kPteExec;

    // 设置用户/内核访问权限
    if (!fIsKernel && (attributes & B_USER_PROTECTION) != 0)
        flags |= kPteUser;

    // 设置缓存属性
    if (memoryType == B_MTR_UC)
        flags |= kPteNoCache;

    // 设置全局标志（内核页面通常是全局的）
    if (fIsKernel)
        flags |= kPteGlobal;

    // 获取页表项
    loongarch_pte *pte = LookupPte(fPageTable, virtualAddress);
    if (pte != NULL && pte->IsValid())
    {
        // 页面已经映射，返回错误
        return B_BAD_VALUE;
    }

    // 分配页表项
    vm_page *pages[4] = {NULL}; // 最多需要4个页表页
    int pageCount = 0;

    loongarch_pte *currentPte = (loongarch_pte *)VirtFromPhys(fPageTable);
    for (int level = 3; level > 0; level--)
    {
        currentPte += VirtAddrPte(virtualAddress, level);
        if (!currentPte->IsValid())
        {
            // 需要分配新的页表页
            vm_page *page = vm_page_allocate_page(reservation, PAGE_STATE_WIRED | VM_PAGE_ALLOC_CLEAR);
            if (page == NULL)
                return B_NO_MEMORY;

            pages[pageCount++] = page;
            phys_addr_t pagePhysical = page->physical_page_number << PAGE_SHIFT;

            // 初始化页表项
            currentPte->val = (pagePhysical & kPtePhysAddrMask) | kPteValid | kPteRead | kPteWrite | kPteExec;

            // 获取下一级页表
            currentPte = (loongarch_pte *)VirtFromPhys(pagePhysical);
        }
        else
        {
            // 使用现有页表
            currentPte = (loongarch_pte *)VirtFromPhys(currentPte->PhysicalAddress());
        }
    }

    // 设置最终的页表项
    currentPte += VirtAddrPte(virtualAddress, 0);
    currentPte->val = (physicalAddress & kPtePhysAddrMask) | flags | kPteValid | kPteAccessed | kPteDirty;

    // 刷新TLB
    asm volatile(
        "invtlb 0, $zero, %0\n"
        :
        : "r"(virtualAddress)
        : "memory");

    return B_OK;
}

// 取消映射内存范围
status_t
LoongArchVMTranslationMap::Unmap(addr_t start, addr_t end)
{
    TRACE("LoongArchVMTranslationMap::Unmap: start %#" B_PRIxADDR ", end %#" B_PRIxADDR "\n", start, end);

    // 确保地址对齐到页边界
    start = ROUNDDOWN(start, B_PAGE_SIZE);
    end = ROUNDUP(end, B_PAGE_SIZE);

    // 遍历范围内的每个页面
    for (addr_t address = start; address < end; address += B_PAGE_SIZE)
    {
        loongarch_pte *pte = LookupPte(fPageTable, address);
        if (pte != NULL && pte->IsValid())
        {
            // 清除页表项
            pte->val = 0;

            // 刷新TLB
            asm volatile(
                "invtlb 0, $zero, %0\n"
                :
                : "r"(address)
                : "memory");
        }
    }

    return B_OK;
}

// 调试标记范围存在
status_t
LoongArchVMTranslationMap::DebugMarkRangePresent(addr_t start, addr_t end, bool markPresent)
{
    // 这个函数用于调试目的，标记一个范围的页面是否存在
    // 在实际实现中，可能需要修改页表项的某些标志位
    return B_OK;
}

// 取消映射单个页面
status_t
LoongArchVMTranslationMap::UnmapPage(VMArea *area, addr_t address, bool updatePageQueue)
{
    TRACE("LoongArchVMTranslationMap::UnmapPage: area %p, address %#" B_PRIxADDR "\n", area, address);

    // 获取页表项
    loongarch_pte *pte = LookupPte(fPageTable, address);
    if (pte == NULL || !pte->IsValid())
        return B_OK; // 页面未映射，不需要操作

    // 获取物理地址
    phys_addr_t physicalAddress = pte->PhysicalAddress();

    // 清除页表项
    pte->val = 0;

    // 刷新TLB
    asm volatile(
        "invtlb 0, $zero, %0\n"
        :
        : "r"(address)
        : "memory");

    // 更新页面队列
    if (updatePageQueue)
    {
        vm_page *page = vm_lookup_page(physicalAddress >> PAGE_SHIFT);
        if (page != NULL)
        {
            DEBUG_PAGE_ACCESS_START(page);
            vm_page_set_state(page, PAGE_STATE_CLEAR);
        }
    }

    return B_OK;
}

// 取消映射多个页面
void LoongArchVMTranslationMap::UnmapPages(VMArea *area, addr_t base, size_t size, bool updatePageQueue)
{
    TRACE("LoongArchVMTranslationMap::UnmapPages: area %p, base %#" B_PRIxADDR ", size %#" B_PRIxSIZE "\n",
          area, base, size);

    // 确保地址对齐到页边界
    base = ROUNDDOWN(base, B_PAGE_SIZE);
    size = ROUNDUP(size, B_PAGE_SIZE);

    // 遍历范围内的每个页面
    for (addr_t address = base; address < base + size; address += B_PAGE_SIZE)
    {
        UnmapPage(area, address, updatePageQueue);
    }
}

// 取消映射整个区域
void LoongArchVMTranslationMap::UnmapArea(VMArea *area, bool deletingAddressSpace, bool ignorePermissions)
{
    TRACE("LoongArchVMTranslationMap::UnmapArea: area %p, base %#" B_PRIxADDR ", size %#" B_PRIxSIZE "\n",
          area, area->Base(), area->Size());

    // 取消映射整个区域的所有页面
    UnmapPages(area, area->Base(), area->Size(), !deletingAddressSpace);
}

// 查询页面信息
status_t
LoongArchVMTranslationMap::Query(addr_t virtualAddress, phys_addr_t *_physicalAddress, uint32 *_flags)
{
    TRACE("LoongArchVMTranslationMap::Query: va %#" B_PRIxADDR "\n", virtualAddress);

    // 获取页表项
    loongarch_pte *pte = LookupPte(fPageTable, virtualAddress);
    if (pte == NULL || !pte->IsValid())
        return B_ENTRY_NOT_FOUND;

    // 获取物理地址
    if (_physicalAddress != NULL)
    {
        *_physicalAddress = pte->PhysicalAddress() + (virtualAddress & (B_PAGE_SIZE - 1));
    }

    // 获取标志
    if (_flags != NULL)
    {
        uint32 flags = 0;

        if (pte->IsReadable())
            flags |= B_READ_AREA;
        if (pte->IsWritable())
            flags |= B_WRITE_AREA;
        if (pte->IsExecutable())
            flags |= B_EXECUTE_AREA;
        if (pte->IsUserAccessible())
            flags |= B_USER_PROTECTION;

        *_flags = flags;
    }

    return B_OK;
}

// 中断查询页面信息
status_t
LoongArchVMTranslationMap::QueryInterrupt(addr_t virtualAddress, phys_addr_t *_physicalAddress, uint32 *_flags)
{
    // 在中断上下文中查询页面信息，通常与普通查询相同
    return Query(virtualAddress, _physicalAddress, _flags);
}

// 保护页面
status_t
LoongArchVMTranslationMap::Protect(addr_t base, addr_t top, uint32 attributes, uint32 memoryType)
{
    TRACE("LoongArchVMTranslationMap::Protect: base %#" B_PRIxADDR ", top %#" B_PRIxADDR ", attributes %#" B_PRIx32 "\n",
          base, top, attributes);

    // 确保地址对齐到页边界
    base = ROUNDDOWN(base, B_PAGE_SIZE);
    top = ROUNDUP(top, B_PAGE_SIZE);

    // 构建页表项标志
    uint64 flags = 0;

    // 设置访问权限
    if ((attributes & B_READ_AREA) != 0)
        flags |= kPteRead;
    if ((attributes & B_WRITE_AREA) != 0 || (attributes & B_KERNEL_WRITE_AREA) != 0)
        flags |= kPteWrite;
    if ((attributes & B_EXECUTE_AREA) != 0)
        flags |= kPteExec;

    // 设置用户/内核访问权限
    if (!fIsKernel && (attributes & B_USER_PROTECTION) != 0)
        flags |= kPteUser;

    // 设置缓存属性
    if (memoryType == B_MTR_UC)
        flags |= kPteNoCache;

    // 设置全局标志（内核页面通常是全局的）
    if (fIsKernel)
        flags |= kPteGlobal;

    // 遍历范围内的每个页面
    for (addr_t address = base; address < top; address += B_PAGE_SIZE)
    {
        loongarch_pte *pte = LookupPte(fPageTable, address);
        if (pte != NULL && pte->IsValid())
        {
            // 保留物理地址和有效位
            phys_addr_t physicalAddress = pte->PhysicalAddress();

            // 更新页表项
            pte->val = (physicalAddress & kPtePhysAddrMask) | flags | kPteValid | kPteAccessed | kPteDirty;

            // 刷新TLB
            asm volatile(
                "invtlb 0, $zero, %0\n"
                :
                : "r"(address)
                : "memory");
        }
    }

    return B_OK;
}

// 清除标志
status_t
LoongArchVMTranslationMap::ClearFlags(addr_t virtualAddress, uint32 flags)
{
    TRACE("LoongArchVMTranslationMap::ClearFlags: va %#" B_PRIxADDR ", flags %#" B_PRIx32 "\n",
          virtualAddress, flags);

    // 获取页表项
    loongarch_pte *pte = LookupPte(fPageTable, virtualAddress);
    if (pte == NULL || !pte->IsValid())
        return B_ENTRY_NOT_FOUND;

    // 清除指定的标志
    uint64 clearMask = 0;
    if ((flags & B_READ_AREA) != 0)
        clearMask |= kPteRead;
    if ((flags & B_WRITE_AREA) != 0)
        clearMask |= kPteWrite;
    if ((flags & B_EXECUTE_AREA) != 0)
        clearMask |= kPteExec;

    // 更新页表项
    pte->val &= ~clearMask;

    // 刷新TLB
    asm volatile(
        "invtlb 0, $zero, %0\n"
        :
        : "r"(virtualAddress)
        : "memory");

    return B_OK;
}

// 清除访问和修改标志
bool LoongArchVMTranslationMap::ClearAccessedAndModified(VMArea *area, addr_t address,
                                                         bool unmapIfUnaccessed, bool &_modified)
{
    TRACE("LoongArchVMTranslationMap::ClearAccessedAndModified: area %p, address %#" B_PRIxADDR "\n",
          area, address);

    // 获取页表项
    loongarch_pte *pte = LookupPte(fPageTable, address);
    if (pte == NULL || !pte->IsValid())
        return false;

    bool accessed = pte->IsAccessed();
    bool modified = pte->IsDirty();

    _modified = modified;

    // 如果页面未被访问且需要取消映射
    if (!accessed && unmapIfUnaccessed)
    {
        UnmapPage(area, address, true);
        return false;
    }

    // 清除访问和修改标志
    if (accessed)
    {
        pte->val &= ~kPteAccessed;
    }
    if (modified)
    {
        pte->val &= ~kPteDirty;
    }

    return true;
}

// 刷新TLB
void LoongArchVMTranslationMap::Flush()
{
    TRACE("LoongArchVMTranslationMap::Flush\n");

    // 刷新整个TLB
    asm volatile(
        "invtlb 0, $zero, $zero\n"
        :
        :
        : "memory");
}

// 更改ASID
void LoongArchVMTranslationMap::ChangeASID(uint32 asid)
{
    TRACE("LoongArchVMTranslationMap::ChangeASID: %u -> %u\n", fASID, asid);

    fASID = asid;

    // 更新CSR寄存器中的ASID
    // 假设ASID存储在CSR寄存器0x18中
    uint64 csr;
    asm volatile(
        "csrrd %0, 0x18\n"
        : "=r"(csr));

    // 清除ASID字段并设置新的ASID
    csr &= ~0xff00ULL; // 假设ASID位于8-15位
    csr |= ((uint64)asid << 8);

    asm volatile(
        "csrwr %0, 0x18\n"
        :
        : "r"(csr)
        : "memory");

    // 刷新TLB
    Flush();
}

// #pragma mark - LoongArchVMPhysicalPageMapper

// 构造函数
LoongArchVMPhysicalPageMapper::LoongArchVMPhysicalPageMapper()
{
    TRACE("LoongArchVMPhysicalPageMapper::ctor\n");
}

// 析构函数
LoongArchVMPhysicalPageMapper::~LoongArchVMPhysicalPageMapper()
{
    TRACE("LoongArchVMPhysicalPageMapper::dtor\n");
}

// 获取物理页面的虚拟地址
status_t
LoongArchVMPhysicalPageMapper::GetPage(phys_addr_t physicalAddress, addr_t *_virtualAddress,
                                       void **_handle)
{
    TRACE("LoongArchVMPhysicalPageMapper::GetPage: pa %#" B_PRIxPHYSADDR "\n", physicalAddress);

    // 使用物理内存映射区域
    *_virtualAddress = VirtFromPhys(physicalAddress);
    *_handle = NULL;

    return B_OK;
}

// 释放物理页面的虚拟地址
status_t
LoongArchVMPhysicalPageMapper::PutPage(addr_t virtualAddress, void *handle)
{
    TRACE("LoongArchVMPhysicalPageMapper::PutPage: va %#" B_PRIxADDR "\n", virtualAddress);

    // 不需要特殊处理
    return B_OK;
}

// 在当前CPU上获取物理页面的虚拟地址
status_t
LoongArchVMPhysicalPageMapper::GetPageCurrentCPU(phys_addr_t physicalAddress,
                                                 addr_t *_virtualAddress, void **_handle)
{
    // 与普通GetPage相同
    return GetPage(physicalAddress, _virtualAddress, _handle);
}

// 在当前CPU上释放物理页面的虚拟地址
status_t
LoongArchVMPhysicalPageMapper::PutPageCurrentCPU(addr_t virtualAddress, void *handle)
{
    // 与普通PutPage相同
    return PutPage(virtualAddress, handle);
}

// 获取用于调试的物理页面的虚拟地址
status_t
LoongArchVMPhysicalPageMapper::GetPageDebug(phys_addr_t physicalAddress,
                                            addr_t *_virtualAddress, void **_handle)
{
    // 与普通GetPage相同
    return GetPage(physicalAddress, _virtualAddress, _handle);
}

// 释放用于调试的物理页面的虚拟地址
status_t
LoongArchVMPhysicalPageMapper::PutPageDebug(addr_t virtualAddress, void *handle)
{
    // 与普通PutPage相同
    return PutPage(virtualAddress, handle);
}

// 设置物理内存
status_t
LoongArchVMPhysicalPageMapper::MemsetPhysical(phys_addr_t address, int value, phys_size_t length)
{
    TRACE("LoongArchVMPhysicalPageMapper::MemsetPhysical: pa %#" B_PRIxPHYSADDR ", length %#" B_PRIxPHYSSIZE "\n",
          address, length);

    // 使用物理内存映射区域
    memset((void *)VirtFromPhys(address), value, length);

    return B_OK;
}

// 从物理内存复制到虚拟内存
status_t
LoongArchVMPhysicalPageMapper::MemcpyFromPhysical(void *to, phys_addr_t from, size_t length,
                                                  bool user)
{
    TRACE("LoongArchVMPhysicalPageMapper::MemcpyFromPhysical: from %#" B_PRIxPHYSADDR ", length %#" B_PRIxSIZE "\n",
          from, length);

    // 使用物理内存映射区域
    if (user)
    {
        // 用户空间复制
        return user_memcpy(to, (const void *)VirtFromPhys(from), length);
    }
    else
    {
        // 内核空间复制
        memcpy(to, (const void *)VirtFromPhys(from), length);
        return B_OK;
    }
}

// 从虚拟内存复制到物理内存
status_t
LoongArchVMPhysicalPageMapper::MemcpyToPhysical(phys_addr_t to, const void *from, size_t length,
                                                bool user)
{
    TRACE("LoongArchVMPhysicalPageMapper::MemcpyToPhysical: to %#" B_PRIxPHYSADDR ", length %#" B_PRIxSIZE "\n",
          to, length);

    // 使用物理内存映射区域
    if (user)
    {
        // 用户空间复制
        return user_memcpy((void *)VirtFromPhys(to), from, length);
    }
    else
    {
        // 内核空间复制
        memcpy((void *)VirtFromPhys(to), from, length);
        return B_OK;
    }
}

// 复制物理页面
void LoongArchVMPhysicalPageMapper::MemcpyPhysicalPage(phys_addr_t to, phys_addr_t from)
{
    TRACE("LoongArchVMPhysicalPageMapper::MemcpyPhysicalPage: to %#" B_PRIxPHYSADDR ", from %#" B_PRIxPHYSADDR "\n",
          to, from);

    // 使用物理内存映射区域
    memcpy((void *)VirtFromPhys(to), (const void *)VirtFromPhys(from), B_PAGE_SIZE);
}